import os, json
from pathlib import Path
from tqdm import tqdm


if __name__ == "__main__":
    native_p = Path(os.getenv("HOME", "")) / "datasets" / "LLaVA-Instruct-150K" / "llava_v1_5_mix665k.json"
    assert native_p.exists(), f"Native LLaVA-Instruct-150K annotation file `{native_p}` must be existed!"
    with open(native_p, "r") as f:
        native_anns = json.load(f)

    only_coco_p = Path(str(native_p).replace(native_p.stem, "llava_v1_5_only_coco_sft"))
    only_coco_anns = []
    for ann in tqdm(native_anns):
        if ('image' in ann) and ('coco' in ann['image']):
            only_coco_anns.append(ann)
    print(f"native llava_v1_5_mix665k file has {len(only_coco_anns)} coco dataset annotations")

    with open(only_coco_p, "w") as f:
        json.dump(only_coco_anns, f)
